import os
import argparse
import xml.etree.ElementTree as ET


CLASSES = ('aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car',
                'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike',
                'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor')


def main(arg):
    info = open('./VOC2007.info', 'w')
    cls = open('./voc.names', 'w')
    for i in CLASSES:
        cls.write(i)
        cls.write('\n')

    for idx, jpg_name in enumerate(os.listdir(arg.img_path)):
        key_name = jpg_name.split('.')[0]
        xml_name = os.path.join(arg.ann_path, key_name + '.xml')
        #parse xml
        tree = ET.parse(xml_name)
        root = tree.getroot()
        size = root.find('size')
        width = size.find('width').text
        height = size.find('height').text
        info.write('{} {} {} {}'.format(idx, os.path.join(arg.img_path, jpg_name), width, height))
        info.write('\n')

        with open('{}/{}'.format(arg.gtp, key_name + '.txt'), 'w') as f:
            for obj in root.iter('object'):
                difficult = int(obj.find('difficult').text)
                cls_name = obj.find('name').text.strip().lower()
                if cls_name not in CLASSES:
                    continue
                xml_box = obj.find('bndbox')
                xmin = (float(xml_box.find('xmin').text))
                ymin = (float(xml_box.find('ymin').text))
                xmax = (float(xml_box.find('xmax').text))
                ymax = (float(xml_box.find('ymax').text))

                if difficult:
                    comment = '{} {} {} {} {} {}'.format(cls_name, xmin, ymin, xmax, ymax, 'difficult')
                else:
                    comment = '{} {} {} {} {}'.format(cls_name, xmin, ymin, xmax, ymax)
                f.write(comment)
                f.write('\n')


def err_msg(msg):
    print('-' * 55)
    print("The specified '{}' file does not exist".format(msg))
    print('You can get the correct parameter information from -h')
    print('-' * 55)
    exit()


def check_args(args):
    if not os.path.exists(args.img_path):
        err_msg(args.img_path)
    if not os.path.exists(args.ann_path):
        err_msg(args.ann_path)
    if not os.path.exists(args.gtp):
        os.makedirs(args.gtp)
    return args


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Parse the VOC2007 dataset label')
    parser.add_argument("--img_path", default="JPEGImages", help='The image path')
    parser.add_argument("--ann_path", default="Annotations", help='Origin xml path')
    parser.add_argument("--gtp", default="ground-truth/", help='The ground true file path')
    args = parser.parse_args()
    args = check_args(args)
    main(args)